package HMM;

import TrainingData.TrainingInstanceCollection;


/**
 * @author Gergely Neu
 * 
 */
public abstract class TrainingMethod {
	
	protected int _iterNr = 0;
	protected TrainingInstanceCollection _trainingData;
	protected TrainingInstanceCollection _testData;
	protected double _emissL2Norm = 0.0;
	protected double _transL2Norm = 0.0;
	protected double _lambda = 0.0; 
	
	/**
	 * Constructor for TrainingMethod.
	 * 
	 * @param iterNr number of iterations to carry out.
	 * @param trainingData the data used for training. 
	 * @param testData the data used for testing.
	 */
	
	protected TrainingMethod(int iterNr, TrainingInstanceCollection trainingData,TrainingInstanceCollection testData){	
		_iterNr = iterNr;
		_trainingData = trainingData;
		_testData = testData;
	}
	
	/**
	 * @author Gergely Neu
	 * 
	 * A pair of sparse matrices used to update the HMM parameters.
	 */
	protected class Updates{
		protected DoubleSparseMatrix _deltaTransMx;
		protected DoubleSparseMatrix _deltaEmissMx;
		
		protected Updates(DoubleSparseMatrix deltaTransMx, DoubleSparseMatrix deltaEmissMx) {
			_deltaTransMx = deltaTransMx;
			_deltaEmissMx = deltaEmissMx;
		}
		
		public DoubleSparseMatrix getDeltaTransMx() {
			return _deltaTransMx;
			}
		
		public DoubleSparseMatrix getDeltaEmissMx() {
			return _deltaEmissMx;
			}
	}
	
	/**
	 * Perform an iteration of the training method. The update itself is not performed, the proposed update is returned instead. 
	 * 
	 * @param transitionMx the transition matrix of the HMM that needs to be trained.
	 * @param emissionMx the emission matrix of the HMM that needs to be trained.
	 * @return the proposed updates.
	 */
	abstract protected Updates iterationStep(DoubleSparseMatrix transitionMx, DoubleSparseMatrix emissionMx);
	
	/**
	 * Computes the step size to be applied at iteration step iter. Normalizes step sizes with the number of training instances to
	 * facilitate step size selection for varying training sets. 
	 * 
	 * @param initialStep the relative size of the first step.
	 * @param iter index of current iteration.
	 * @return the step size to be used.
	 */
	protected double computeStepSize(double initialStep, int iter) {
		return initialStep / Math.sqrt(iter+1.0) / _trainingData.getSize(); //TODO need to implement more step size choices -- G.
		//return initialStep / _trainingData.getSize();
	}
	
	/**
	 * Trains the hidden Markov model specified by the matrices transitionMx and emissionMx iteratively. The HMM is tested on the training
	 * and test data sets on each iteration.
	 * 
	 * @param transitionMx the transition matrix of the HMM to train
	 * @param emissionMx the emission matrix of the HMM to train
	 * 
	 * @see HMM
	 * 
	 */
	
	public void train(DoubleSparseMatrix transitionMx, DoubleSparseMatrix emissionMx){	
		
		for (int i = 0; i < _iterNr; i++){
			System.out.println("[Iteration " + i + "]");

			HMM testHMM = new HMM(transitionMx, emissionMx);
			
			System.out.print("[test]    ");
			testHMM.test(_testData, System.out, false);
			
			System.out.print("[training]");
			testHMM.test(_trainingData, System.out, false);
			
			Updates updates = iterationStep(transitionMx,emissionMx);
			
			double stepSize = computeStepSize(1e5, i);
			
			if (i == 0) {
				_emissL2Norm = updates.getDeltaEmissMx().getL2Norm();
				_transL2Norm = updates.getDeltaTransMx().getL2Norm();
			}
			
			_lambda = 0.0;
			
			transitionMx.scalarTimes(1-_lambda);
			emissionMx.scalarTimes(1-_lambda);
			
			transitionMx.add(updates.getDeltaTransMx().scalarTimes(stepSize / _transL2Norm));
			emissionMx.add(updates.getDeltaEmissMx().scalarTimes(stepSize / _emissL2Norm));
		}
	
	}
}
